package com.example.springai.controller;

import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingOptionsBuilder;
import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.model.Media;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.ai.ollama.api.OllamaModel;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;
import org.springframework.core.io.FileSystemResource;
import org.springframework.core.io.Resource;
import org.springframework.util.MimeTypeUtils;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.Arrays;
import java.util.List;
import java.util.Map;

@RestController
public class OllamaController {

    @Autowired
    private OllamaChatModel chatModel;

    @Autowired
    private OllamaEmbeddingModel ollamaEmbeddingModel;

    @Autowired
    private VectorStore vectorStore;

    @GetMapping(value = "/ai/generateStream", produces = "text/html;charset=UTF-8")
    public Flux<String> generateStream(@RequestParam(value = "message", defaultValue = "给我讲一个笑话") String message) {
        List<Message> messageList = Arrays.asList(new SystemMessage("请全部用简体中文回答"), new UserMessage(message));
        Prompt prompt = new Prompt(messageList);
        return chatModel.stream(prompt).map((r) -> r.getResult() != null && r.getResult().getOutput() != null && r.getResult().getOutput().getContent() != null ? r.getResult().getOutput().getContent() : "").filter(StringUtils::hasLength);
    }

    @GetMapping(value = "/ai/generateEmbedding")
    public String generateEmbedding(@RequestParam(value = "message", defaultValue = "给我讲一个笑话") String message) {
        EmbeddingRequest request = new EmbeddingRequest(List.of(message), EmbeddingOptionsBuilder.builder().build());
        EmbeddingResponse embeddingResponse = ollamaEmbeddingModel.call(request);
        return Arrays.toString(embeddingResponse.getResult().getOutput());
    }

    @GetMapping(value = "/ai/txtToVectorstore")
    public String txtToVectorstore() {
        //从本地txt文件读取文本内容
        Resource resource = new FileSystemResource("C:\\Users\\yurii\\Desktop\\JAVA-RAG.txt");
        List<Document> documents = new TextReader(resource).read();
        //文本内容分割
        List<Document> splitDocuments = new TokenTextSplitter(220, 220, 5, 100, true).transform(documents);
        //文本内容向量化并存储
        vectorStore.add(splitDocuments);
        return "ok";
    }

    @GetMapping(value = "/ai/generateFromVectorstore")
    public String generateFromVectorstore(@RequestParam(value = "message", defaultValue = "Spring") String message) {
        //向量数据检索
        List<Document> results = vectorStore.similaritySearch(SearchRequest.query(message).withTopK(5));
        return results.get(0).getContent();
    }

    @GetMapping(value = "/ai/generateWithPromptTemplate", produces = "text/html;charset=UTF-8")
    public Flux<String> generateWithPromptTemplate(@RequestParam(value = "singer", defaultValue = "周杰伦") String singer) {
//        String template = "请问{singer}最受欢迎的歌有哪些？";
        Resource resource = new FileSystemResource("C:\\Users\\yurii\\Desktop\\template.txt");
        PromptTemplate promptTemplate = new PromptTemplate(resource);
        // 动态地将singer填充进去
        Prompt prompt = promptTemplate.create(Map.of("singer", singer));
        return chatModel.stream(prompt).map((r) -> r.getResult() != null && r.getResult().getOutput() != null && r.getResult().getOutput().getContent() != null ? r.getResult().getOutput().getContent() : "").filter(StringUtils::hasLength);
    }

    @GetMapping(value = "/ai/javaAssistantWithRAG", produces = "text/html;charset=UTF-8")
    public Flux<String> javaAssistantWithRAG(@RequestParam(value = "message", defaultValue = "synchronized和volatile有什么区别？") String message) {
        List<Document> results = vectorStore.similaritySearch(SearchRequest.query(message).withTopK(5));
        String template = "你需要使用文档内容对用户提出的问题进行回复，同时你需要表现得天生就知道这些内容，" + "不能在回复中体现出你是根据给出的文档内容进行回复的，这点非常重要。" + "当用户提出的问题无法根据文档内容进行回复或者你也不清楚时，回复不知道即可。" + "文档内容如下：" + "{information}";
        PromptTemplate promptTemplate = new PromptTemplate(template);
        Prompt prompt = promptTemplate.create(Map.of("information", results.get(0).getContent()));
        return chatModel.stream(prompt).map((r) -> r.getResult() != null && r.getResult().getOutput() != null && r.getResult().getOutput().getContent() != null ? r.getResult().getOutput().getContent() : "").filter(StringUtils::hasLength);
    }

    @GetMapping(value = "/ai/generateFromMultimodality", produces = "text/html;charset=UTF-8")
    public Flux<String> generateFromMultimodality(@RequestParam(value = "message", defaultValue = "请问图片里有什么？") String message) {
        Resource imageResource = new FileSystemResource("C:\\Users\\yurii\\Desktop\\multimodality-test.png");
        List<Message> messageList = Arrays.asList(new SystemMessage("请全部用简体中文回答"),
                new UserMessage(message, new Media(MimeTypeUtils.IMAGE_PNG, imageResource)));
        return chatModel.stream(new Prompt(messageList, OllamaOptions.builder().withModel(OllamaModel.LLAVA).build()))
                .map((r) -> r.getResult() != null && r.getResult().getOutput() != null && r.getResult().getOutput().getContent() != null ? r.getResult().getOutput().getContent() : "")
                .filter(StringUtils::hasLength);
    }

}
